
/**
  ******************************************************************************
  * Copyright 2021 The Microbee Authors. All Rights Reserved.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  * 
  * http://www.apache.org/licenses/LICENSE-2.0
  * 
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * 
  * @file       mode_acro.c
  * @author     baiyang
  * @date       2023-7-13
  ******************************************************************************
  */

/*----------------------------------include-----------------------------------*/
#include "mode.h"
#include "fms.h"
/*-----------------------------------macro------------------------------------*/

/*----------------------------------typedef-----------------------------------*/

/*---------------------------------prototype----------------------------------*/
static bool macro_enter(mode_base_t mode);
static enum ModeNumber macro_number(mode_base_const_t mode);
static const char *macro_name(mode_base_const_t mode);
static const char *macro_name4(mode_base_const_t mode);
static void macro_update(mode_base_t mode);
static void macro_run(mode_base_t mode);
static void macro_stabilize(mode_base_t mode);
static void macro_stabilize_quaternion(mode_base_t mode);
/*----------------------------------variable----------------------------------*/
static struct mode_ops mode_acro_ops = {
        .run                                = macro_run,
        .mode_number                        = macro_number,
        .name                               = macro_name,
        .name4                              = macro_name4,
        .update                             = macro_update,
        .is_vtol_mode                       = NULL,
        .is_vtol_man_throttle               = NULL,
        .is_vtol_man_mode                   = NULL,
        .is_guided_mode                     = NULL,
        .allows_terrain_disable             = NULL,
        .does_automatic_thermal_switch      = NULL,
        .navigate                           = NULL,
        .allows_throttle_nudging            = NULL,
        .does_auto_navigation               = NULL,
        .does_auto_throttle                 = NULL,
        .mode_allows_autotuning             = NULL,
        .update_target_altitude             = NULL,
        .handle_guided_request              = NULL,
        .is_landing                         = NULL,
        .is_taking_off                      = NULL,
        ._enter                             = macro_enter,
        ._exit                              = NULL,
        ._pre_arm_checks                    = NULL};
/*-------------------------------------os-------------------------------------*/

/*----------------------------------function----------------------------------*/
void macro_ctor(struct ModeAcro *mode_acro)
{
    mode_ctor(&mode_acro->mode, &mode_acro_ops);
}

static bool macro_enter(mode_base_t mode)
{
    struct ModeAcro *mode_acro = (struct ModeAcro *)mode;

    mode_acro->acro_state.locked_roll = false;
    mode_acro->acro_state.locked_pitch = false;
    mode_acro->acro_state.q = fms.ahrs->attitude_body;
    return true;
}

// returns a unique number specific to this mode
static enum ModeNumber macro_number(mode_base_const_t mode)
{
    (void)mode;
    return MN_ACRO;
}

// returns full text name
static const char *macro_name(mode_base_const_t mode)
{
    (void)mode;
    return "ACRO";
}

// returns a string for this flightmode, exactly 4 bytes
static const char *macro_name4(mode_base_const_t mode)
{
    (void)mode;
    return "ACRO";
}

static void macro_update(mode_base_t mode)
{
    struct ModeAcro *mode_acro = (struct ModeAcro *)mode;

    // handle locked/unlocked control
    if (mode_acro->acro_state.locked_roll) {
        fms.nav_roll_cd = mode_acro->acro_state.locked_roll_err;
    } else {
        fms.nav_roll_cd = fms.ahrs->roll_sensor_cd;
    }
    if (mode_acro->acro_state.locked_pitch) {
        fms.nav_pitch_cd = mode_acro->acro_state.locked_pitch_cd;
    } else {
        fms.nav_pitch_cd = fms.ahrs->pitch_sensor_cd;
    }
}

static void macro_run(mode_base_t mode)
{
    if (fms.g.acro_locking == 2 && fms.g.acro_yaw_rate > 0 &&
        yctrl_rate_control_enabled(&fms.yaw_controller)) {
        // we can do 3D acro locking
        macro_stabilize_quaternion(mode);
        return;
    }

    // Normal acro
    macro_stabilize(mode);
}

/*
  this is the ACRO mode stabilization function. It does rate
  stabilization on roll and pitch axes
 */
static void macro_stabilize(mode_base_t mode)
{
    struct ModeAcro *mode_acro = (struct ModeAcro *)mode;

    const float speed_scaler = fms.surface_speed_scaler;
    const float rexpo = fms_roll_in_expo(true);
    const float pexpo = fms_pitch_in_expo(true);
    float roll_rate = (rexpo/SERVO_MAX) * fms.g.acro_roll_rate;
    float pitch_rate = (pexpo/SERVO_MAX) * fms.g.acro_pitch_rate;

    mode_acro->acro_state.q = fms.ahrs->attitude_body;;

    /*
      check for special roll handling near the pitch poles
     */
    if (fms.g.acro_locking && math_flt_zero(roll_rate)) {
        /*
          we have no roll stick input, so we will enter "roll locked"
          mode, and hold the roll we had when the stick was released
         */
        if (!mode_acro->acro_state.locked_roll) {
            mode_acro->acro_state.locked_roll = true;
            mode_acro->acro_state.locked_roll_err = 0;
        } else {
            mode_acro->acro_state.locked_roll_err += fms.ahrs->gyr.x * fms.G_Dt;
        }
        int32_t roll_error_cd = -ToDeg(mode_acro->acro_state.locked_roll_err)*100;
        fms.nav_roll_cd = fms.ahrs->roll_sensor_cd + roll_error_cd;
        // try to reduce the integrated angular error to zero. We set
        // 'stabilze' to true, which disables the roll integrator
        srv_channels_set_output_scaled(k_aileron, rctrl_get_servo_out(&fms.roll_controller, roll_error_cd,
                                                                                             speed_scaler,
                                                                                             true, false));
    } else {
        /*
          aileron stick is non-zero, use pure rate control until the
          user releases the stick
         */
        mode_acro->acro_state.locked_roll = false;
        srv_channels_set_output_scaled(k_aileron, rctrl_get_rate_out(&fms.roll_controller, roll_rate,  speed_scaler));
    }

    if (fms.g.acro_locking && math_flt_zero(pitch_rate)) {
        /*
          user has zero pitch stick input, so we lock pitch at the
          point they release the stick
         */
        if (!mode_acro->acro_state.locked_pitch) {
            mode_acro->acro_state.locked_pitch = true;
            mode_acro->acro_state.locked_pitch_cd = fms.ahrs->pitch_sensor_cd;
        }
        // try to hold the locked pitch. Note that we have the pitch
        // integrator enabled, which helps with inverted flight
        fms.nav_pitch_cd = mode_acro->acro_state.locked_pitch_cd;
        srv_channels_set_output_scaled(k_elevator, pctrl_get_servo_out(&fms.pitch_controller, fms.nav_pitch_cd - fms.ahrs->pitch_sensor_cd,
                                                                                               speed_scaler,
                                                                                               false, false));
    } else {
        /*
          user has non-zero pitch input, use a pure rate controller
         */
        mode_acro->acro_state.locked_pitch = false;
        srv_channels_set_output_scaled(k_elevator, pctrl_get_rate_out(&fms.pitch_controller, pitch_rate, speed_scaler));
    }

    float rudder_output;
    if (fms.g.acro_yaw_rate > 0 && yctrl_rate_control_enabled(&fms.yaw_controller)) {
        // user has asked for yaw rate control with yaw rate scaled by ACRO_YAW_RATE
        const float rudd_expo = fms_rudder_in_expo(true);
        const float yaw_rate = (rudd_expo/SERVO_MAX) * fms.g.acro_yaw_rate;
        rudder_output = yctrl_get_rate_out(&fms.yaw_controller , yaw_rate,  speed_scaler, false);
    } else if (fms_flight_option_enabled(FO_ACRO_YAW_DAMPER)) {
        // use yaw controller
        rudder_output = (float)fms_calc_nav_yaw_coordinated();
    } else {
        /*
          manual rudder
        */
        rudder_output = (float)fms_rudder_input();
    }

    mode_output_rudder_and_steering(rudder_output);
}

/*
  quaternion based acro stabilization with continuous locking. Enabled with ACRO_LOCKING=2
 */
static void macro_stabilize_quaternion(mode_base_t mode)
{
    struct ModeAcro *mode_acro = (struct ModeAcro *)mode;

    const float speed_scaler = fms.surface_speed_scaler;
    Quat_t *q = &mode_acro->acro_state.q;
    const float rexpo = fms_roll_in_expo(true);
    const float pexpo = fms_pitch_in_expo(true);
    const float yexpo = fms_rudder_in_expo(true);

    // get pilot desired rates
    float roll_rate = (rexpo/SERVO_MAX) * fms.g.acro_roll_rate;
    float pitch_rate = (pexpo/SERVO_MAX) * fms.g.acro_pitch_rate;
    float yaw_rate = (yexpo/SERVO_MAX) * fms.g.acro_yaw_rate;
    bool roll_active = !math_flt_zero(roll_rate);
    bool pitch_active = !math_flt_zero(pitch_rate);
    bool yaw_active = !math_flt_zero(yaw_rate);

    // integrate target attitude
    Vector3f_t r = { radians(roll_rate), radians(pitch_rate), radians(yaw_rate) };
    r.x *= fms.G_Dt;
    r.y *= fms.G_Dt;
    r.z *= fms.G_Dt;

    quat_rotate_fast(q, &r);
    quat_norm(q);

    // fill in target roll/pitch for GCS/logs
    fms.nav_roll_cd = degrees(quat_get_roll(q))*100;
    fms.nav_pitch_cd = degrees(quat_get_pitch(q))*100;

    // get AHRS attitude
    Quat_t ahrs_q = fms.ahrs->attitude_body;;

    // zero target if not flying, no stick input and zero throttle
    if (math_flt_zero(fms_get_throttle_input(false)) &&
        !fms_is_flying() &&
        math_flt_zero(roll_rate) &&
        math_flt_zero(pitch_rate) &&
        math_flt_zero(yaw_rate)) {
        // cope with sitting on the ground with neutral sticks, no throttle
        *q = ahrs_q;
    }

    // get error in attitude
    Quat_t ahrs_q_inv;
    quat_inverse(&ahrs_q_inv, &ahrs_q);

    Quat_t error_quat = quat_mult(&ahrs_q_inv, q);
    Vector3f_t error_angle1;
    quat_to_axis_angle(&error_quat, &error_angle1);

    // don't let too much error build up, limit to 0.2s
    const float max_error_t = 0.2f;
    float max_err_roll_rad  = radians(fms.g.acro_roll_rate*max_error_t);
    float max_err_pitch_rad = radians(fms.g.acro_pitch_rate*max_error_t);
    float max_err_yaw_rad   = radians(fms.g.acro_yaw_rate*max_error_t);

    if (!roll_active && mode_acro->acro_state.roll_active_last) {
        max_err_roll_rad = 0;
    }
    if (!pitch_active && mode_acro->acro_state.pitch_active_last) {
        max_err_pitch_rad = 0;
    }
    if (!yaw_active && mode_acro->acro_state.yaw_active_last) {
        max_err_yaw_rad = 0;
    }

    Vector3f_t desired_rates = error_angle1;
    desired_rates.x = math_constrain_float(desired_rates.x, -max_err_roll_rad, max_err_roll_rad);
    desired_rates.y = math_constrain_float(desired_rates.y, -max_err_pitch_rad, max_err_pitch_rad);
    desired_rates.z = math_constrain_float(desired_rates.z, -max_err_yaw_rad, max_err_yaw_rad);

    // correct target based on max error
    Vector3f_t desired_rates2 = {desired_rates.x - error_angle1.x,
                                 desired_rates.y - error_angle1.y,
                                 desired_rates.z - error_angle1.z};

    quat_rotate_fast(q, &desired_rates2);
    quat_norm(q);

    // convert to desired body rates
    desired_rates.x /= rctrl_get_tau(&fms.roll_controller);
    desired_rates.y /= pctrl_get_tau(&fms.pitch_controller);
    desired_rates.z /= pctrl_get_tau(&fms.pitch_controller); // no yaw tau parameter, use pitch

    desired_rates.x *= degrees(1.0f);
    desired_rates.y *= degrees(1.0f);
    desired_rates.z *= degrees(1.0f);

    if (roll_active) {
        desired_rates.x = roll_rate;
    }
    if (pitch_active) {
        desired_rates.y = pitch_rate;
    }
    if (yaw_active) {
        desired_rates.z = yaw_rate;
    }

    // call to rate controllers
    srv_channels_set_output_scaled(k_aileron, rctrl_get_rate_out(&fms.roll_controller, desired_rates.x, speed_scaler));
    srv_channels_set_output_scaled(k_elevator, pctrl_get_rate_out(&fms.pitch_controller, desired_rates.y, speed_scaler));
    mode_output_rudder_and_steering(yctrl_get_rate_out(&fms.yaw_controller, desired_rates.z,  speed_scaler, false));

    mode_acro->acro_state.roll_active_last = roll_active;
    mode_acro->acro_state.pitch_active_last = pitch_active;
    mode_acro->acro_state.yaw_active_last = yaw_active;
}
/*------------------------------------test------------------------------------*/


